import torch
import numpy as np
import pandas as pd
import os
from PIL import Image
import re
import random

def integers_to_onehot(integers, max_value=10):
    # 创建一个全为 0 的矩阵，行数为整数的个数，列数为 max_value + 1
    integers = [min(i, max_value) for i in integers]
    onehot_matrix = torch.zeros((len(integers), max_value + 1))
    # 为每个整数设置 one-hot 编码
    for i, integer in enumerate(integers):
        onehot_matrix[i, integer] = 1
    return onehot_matrix

def get_labels(path):
    data = pd.read_excel(path, index_col=0)
    data = data.dropna(axis=0, how='any')
    labels, convert_dict = {}, {}

    for i in range(5):
        convert_dict[i] = i * 0.15 + 0.2
        # convert_dict[i] = i
    first_name = data.iloc[0].tolist()
    for name in data[1:].index.values:
        labels[name] = data.loc[name].values.tolist()[:-1]
        # temp = data.loc[name].values.tolist()[:-1]
        # labels[name] = temp
        temp = []
        for i in labels[name]:
            temp.append(int(i))
        labels[name] = temp
        labels[name] = [convert_dict[min(int(i), 4)] for i in labels[name]]
    return labels

def get_last_label(path):
    data = pd.read_excel(path, index_col=0)
    data = data.dropna(axis=0, how='any')
    labels = {}
    for name in data[1:].index.values:
        # labels[name] = int(data.loc[name].values.tolist()[-1])
        labels[name] = int(data.loc[name].values.tolist()[-1])
        labels[name] = integers_to_onehot([labels[name]], max_value=10)

    return labels

def extract_string(string):
    pattern = r'《(.*?)》'
    matches = re.findall(pattern, string)
    return matches

def get_paint_to_course(paint_to_course, path):
    df = pd.read_excel(path, index_col=0, skiprows=1)
    df = df.fillna("empty")
    for index, row in df.iterrows():
        if row[0] == 'empty':
            continue
        paint_name, course_name = index, row[0]
        course_name = extract_string(course_name)[0]
        paint_to_course[paint_name] = course_name
    return paint_to_course

def get_course_to_dimension(course_to_dimension, path):
    df = pd.read_excel(path)
    df = df.fillna("empty")
    for index, row in df.iterrows():
        if row[0] != "empty":
            course_name = row[0]
            course_name = extract_string(course_name)[0]
            course_to_dimension[course_name] = [row[3].replace(' ', '')]
        elif row[0] == row[1] and row[0] == row[2] and row[0] == row[3] and row[0] == "empty":
            break
        else:
            course_to_dimension[course_name].append(row[3].replace(' ', ''))
    return course_to_dimension

def get_dimension_to_idx(dimension_to_idx, path):
    data = pd.read_excel(path, index_col=0)
    data = data.dropna(axis=0, how='any')
    column_names = data.iloc[0]
    for idx in range(len(column_names)):
        column_names[idx] = column_names[idx].replace(' ', '')
        dimension_to_idx[column_names[idx]] = idx
    return dimension_to_idx

def get_pbs_dimension(path):
    # print(path)
    data = pd.read_excel(path, header=None)
    Data = pd.read_excel(path)
    temp = Data.head()
    data = data.fillna("empty")
    pbs_dimension, row_idx = {}, 2
    max_min_scaler = lambda x : (x) / (7)
    col_name = data.iloc[1, 3:-1].tolist()
    while row_idx < data.shape[0]:
        name = str(data.loc[row_idx][0])
        if name == "empty":
            break

        pb_counts = int(data.loc[row_idx][1])
        # print(name)
        pb_df = data.iloc[row_idx : row_idx + pb_counts - 1, 3 : -1].astype(int)
        for row in range(len(pb_df)):
            for col in range(0, 77, 1):
                # print(pb_df.iloc[row, col])
                if pb_df.iloc[row, col] >= 8:
                    te = pb_df.iloc[row, col]
                    with open("data_error.txt", mode='a') as file:
                        ans = path + str(' ') + name + str(' ') + str(row + 1) + str(' ') + str(col_name[col]) + str('\n')
                        file.write(ans)
        # print('yes')
        # pb_df = pb_df.apply(max_min_scaler)
        pbs_dimension[name] = [sum(pb_df.loc[ : , _]) / pb_counts for _ in range(3, len(data.columns) - 1, 1)]
        pbs_dimension[name] = [max_min_scaler(_) for _ in pbs_dimension[name]]
        row_idx = row_idx + pb_counts
    return pbs_dimension

def get_all_pbs_dimension(path):
    pbs_dimension, pbs_info_idx, pbs_nd_update = {}, {}, {}
    for filename in os.listdir((path)):
        deep_path = os.path.join(path, filename)
        pb = get_pbs_dimension(deep_path)
        pbs_dimension.update(pb)
    now_id = 0
    for ke, va in pbs_dimension.items():
        pbs_info_idx[now_id] = ke
        now_id += 1
        great_D = []
        for _ in range(len(va)):
            if va[_] >= 0.7:
                great_D.append(_)
        if len(great_D) >= 5 and len(great_D) <= 20:
            pbs_nd_update[ke] = great_D
        else:
            b = sorted(enumerate(va), key=lambda x : x[1])
            pbs_nd_update[ke] = [_[0] for _ in b[:10]]
        
    return pbs_dimension, pbs_info_idx, pbs_nd_update

def get_all_reflect_relation(data_path):
    path = f'{data_path}/打分表'
    paint_to_course, dimension_to_idx = {}, {}
    for filename in os.listdir((path)):
        deep_path = os.path.join(path, filename)
        if '表' in filename:
            paint_to_course = get_paint_to_course(paint_to_course, deep_path)
    dimension_to_idx = get_dimension_to_idx(dimension_to_idx, f'{data_path}/画/打分1.xlsx')

    DF = pd.read_excel(f"{data_path}/1.xlsx", header=None)
    DF = DF.fillna("empty")
    first_row = DF.iloc[0, :-1].tolist()
    second_row = DF.iloc[1, :-1].tolist()

    ability = {}
    pos = 0
    while pos < len(first_row):
        while pos < len(first_row) and first_row[pos] == str('empty'):
            pos += 1
        idx = pos + 1
        ability[first_row[pos]] = []
        ability[first_row[pos]].append(second_row[pos].replace(" ", ""))
        while idx < len(first_row) and first_row[idx] == str('empty'):
            ability[first_row[pos]].append(second_row[idx].replace(" ", ""))
            idx += 1
        pos = idx

    dimension_adj_matrix = torch.zeros(77, 77).cuda()
    for ke, va in ability.items():
        for i in range(len(va)):
            for j in range(i + 1, len(va), 1):
                i_idx, j_idx = dimension_to_idx[va[i]], dimension_to_idx[va[j]]
                dimension_adj_matrix[i_idx, j_idx], dimension_adj_matrix[j_idx, i_idx] = 1, 1
    return paint_to_course, dimension_adj_matrix

# def find_dir_files(path):
#     for filename in os.listdir((path)):
#         deep_path = os.path.join(path, filename)
#         if filename.endswith('.xlsx'):

def get_RL_data(image_size, data_path):
    image_path = f'{data_path}/画1'
    image_list, labels, last_label = {}, {}, {}
    paint_to_course, _  = get_all_reflect_relation(data_path)

    for filename in os.listdir((image_path)):
        deep_path = os.path.join(image_path, filename)
        if filename.endswith('.xlsx'):
            labels.update(get_labels(deep_path))
            last_label.update(get_last_label(deep_path))
        if not os.path.isdir(deep_path):
            continue
        for wholefilename in os.listdir((deep_path)):
            if wholefilename.endswith('.jpg') or wholefilename.endswith('.png') or wholefilename.endswith('.jpeg'):
                img_path = os.path.join(deep_path, wholefilename)
                img = Image.open(img_path).convert('RGB')
                img = img.resize((image_size, image_size))
                img = torch.tensor(np.array(img)).permute(2, 0, 1).float()
                name_without_extension = os.path.splitext(wholefilename)[0]
                if name_without_extension in paint_to_course.keys():
                    image_list[name_without_extension] = img
    image_list_temp = {}
    for ke in (image_list.keys()):
        if ke not in labels or ke not in last_label:
            continue
        va = last_label[ke][0].tolist()
        idx = va.index(max(va))
        image_list_temp[ke] = image_list[ke]
    image_list = image_list_temp

    train_set, test_set, cnt, tl = [], [], 0, len(image_list)
    for ke, va in image_list.items():
        cnt += 1
        if cnt <= int(tl * 0.7):
            train_set.append([va, ke, labels[ke]])
        else:
            test_set.append([va, ke, labels[ke]])
    return train_set, test_set

def get_data(image_size):
    image_path = './画'
    image_list, labels, last_label = {}, {}, {}
    paint_to_course, course_to_dimension, dimension_to_idx = get_all_reflect_relation()

    for filename in os.listdir((image_path)):
        deep_path = os.path.join(image_path, filename)
        if filename.endswith('.xlsx'):
            labels.update(get_labels(deep_path))
            last_label.update(get_last_label(deep_path))
        if not os.path.isdir(deep_path):
            continue
        for wholefilename in os.listdir((deep_path)):
            if wholefilename.endswith('.jpg') or wholefilename.endswith('.png') or wholefilename.endswith('.jpeg'):
                img_path = os.path.join(deep_path, wholefilename)
                img = Image.open(img_path).convert('RGB')
                img = img.resize((image_size, image_size))
                img = torch.tensor(np.array(img)).permute(2, 0, 1).float()
                name_without_extension = os.path.splitext(wholefilename)[0]
                if name_without_extension in paint_to_course.keys():
                    image_list[name_without_extension] = img
    image_list_temp = {}
    for ke in (image_list.keys()):
        if ke not in labels or ke not in last_label:
            continue
        va = last_label[ke][0].tolist()
        idx = va.index(max(va))
        image_list_temp[ke] = image_list[ke]
    image_list = image_list_temp

    data = pd.read_excel('./画/打分量表2 于.xlsx', index_col=0)
    data = data.dropna(axis=0, how='any')
    first_name = data.iloc[0].tolist()
    Data = []
    DData = []
    for ke, va in labels.items():
        if ke not in labels or ke not in last_label or ke not in paint_to_course:
            continue
        temp = []
        temp.append(ke)
        for idx in range(len(first_name) - 1):
            temp.append(va[idx])
        temp.append(last_label[ke])
        Data.append(temp)

        ttemp = []
        ttemp.append(ke)
        arr = []
        Temp = course_to_dimension[paint_to_course[ke]]
        for idx in range(len(Temp)):
            arr.append([dimension_to_idx[Temp[idx]], va[dimension_to_idx[Temp[idx]]]])
        sort_arr = sorted(arr, key=lambda x: x[1])
        ttemp.append(first_name[sort_arr[0][0]])
        DData.append(ttemp)


    Row_name = ['Name']
    Row_name = Row_name + first_name
    # wdf = pd.DataFrame(Data, columns=Row_name)
    # wdf.to_excel('/home/s44/AestheticPerception_v2/ND.xlsx', index=False, header=True, engine='openpyxl')
    # wwdf = pd.DataFrame(DData, columns=['Name', 'min'])
    # wwdf.to_excel('/home/s44/AestheticPerception_v2/ND1.xlsx', index=False, header=True, engine='openpyxl')
    train_set, test_set, valid_set = [], [], []
    cnt = 0
    image_list_temp = {}
    train_dict, data_dict, valid_dict = {}, {}, {}
    for i in range(11):
        train_dict[i], data_dict[i], valid_dict[i] = [], 0, 0

    for ke in (image_list.keys()):
        if ke not in labels or ke not in last_label:
            continue
        va = last_label[ke][0].tolist()
        idx = va.index(max(va))
        image_list_temp[ke] = image_list[ke]
    image_list = image_list_temp

    student_cnt = {}
    for ke in (image_list.keys()):
        pattern = re.compile(r'[\u4e00-\u9fff]+')
        student_name = pattern.findall(ke)
        if student_name[0] not in student_cnt:
            student_cnt[student_name[0]] = 0
        student_cnt[student_name[0]] += 1
        cnt += 1
        if cnt <= (int)(len(image_list) * 0.7):
            train_set.append([image_list[ke], labels[ke], last_label[ke], ke])
            
        elif cnt <= (int)(len(image_list) * 0.8):
            valid_set.append(([image_list[ke], labels[ke], last_label[ke]], ke))
        else:
            test_set.append([image_list[ke], labels[ke], last_label[ke], ke])
    # print(f'student num is {len(student_cnt)}, image num is {len(image_list)}')
    return train_set, test_set, train_dict, valid_set
def collate(data):
    img = [i[0].tolist() for i in data]
    labels = []
    for i in data:
        one_batch = []
        for j in i[1]:
            one_batch.append(j)
        labels.append(one_batch)
    last_label = []
    for i in data:
        one_batch = i[2].tolist()
        last_label.append(one_batch)
    return torch.tensor(img), torch.tensor(labels), torch.tensor(last_label)

def collate_RL(data):
    img = [i[0].tolist() for i in data]
    name = [i[1] for i in data]
    labels = []
    for i in data:
        one_batch = []
        for j in i[-1]:
            # one_batch.append(j.item())
            one_batch.append(j)
        labels.append(one_batch)
    return torch.tensor(img), name, torch.tensor(labels)

def collate_test(data):
    img = [i[0].tolist() for i in data]
    labels = []
    for i in data:
        one_batch = []
        for j in i[1]:
            # one_batch.append(j.item())
            one_batch.append(j)
        labels.append(one_batch)
    last_label = []
    for i in data:
        one_batch = i[2].tolist()
        # last_label.append(one_batch[0])
        last_label.append(one_batch)
    name = [i[3] for i in data]
    is_mood = [i[4] for i in data]
    return torch.tensor(img), torch.tensor(labels), torch.tensor(last_label), name, is_mood

def collate_v2(data):
    img = [i[0].tolist() for i in data]
    labels = []
    for i in data:
        one_batch = []
        for j in i[1]:
            one_batch.append(j)
        labels.append(one_batch)
    last_label = []
    for i in data:
        one_batch = i[2].tolist()
        last_label.append(one_batch)
    name = [i[3] for i in data]
    return torch.tensor(img), torch.tensor(labels), torch.tensor(last_label), name